import UserCf
import argparse
import pandas as pd
from Experiment import Experiment
import ItemCf
import dataframe_image as dfi
import os

parser = argparse.ArgumentParser(description="user cf recommendation")
parser.add_argument('--train_rate', default=0.9, type=float,
                    help="This is the training data's rate  ")
parser.add_argument('--epoch', default=1, type=int
                    , help="how many epoch it takes")
parser.add_argument('--model', default="user_cf", type=str,
                    help="Which model used to predict,it only can be "
                         "'user_cf or item_cf'")
parser.add_argument('--K', default=5, type=int, help="Top K people in matrix will be considered")
parser.add_argument('--N', default=10, type=int, help="Automatically choose N items to recommend")
parser.add_argument('--path',default='./output.txt',help="output file path default is the home directory of this "
                                                         "project")
args = parser.parse_args()

if args.model == "user_cf":
    model = UserCf.UserCf(output_path=args.path)
elif args.model == "item_cf":
    model = ItemCf.ItemCf(output_path=args.path)
else:
    raise AssertionError("model can only be 'user_cf' or 'item_cf' ")

print("train_rate : " + str(args.train_rate))
print("epoch :" + str(args.epoch))
print("model :" + str(args.model))
print("K :" + str(args.K))
print("N :"+str(args.N))
print("output file path :"+str(args.path))

if os.path.exists(args.path):
    os.remove(args.path)
if os.path.exists('./tmp.txt'):
    os.remove('./tmp.txt')

exper = Experiment(K=args.K, N=args.N, epoch=args.epoch, train_rate=args.train_rate, model=model)
output = exper.run()


o = pd.DataFrame(
    output,
    columns=["Precision", "Recall", "Coverage", "Popularity"]
)
mean = pd.DataFrame(o.mean())
mean.columns = ['mean']
output_with_mean = pd.concat([o, mean.transpose()], axis=0)
print(output_with_mean)
# 保存输出图片
dfi.export(output_with_mean,"./output.png")
print("save the output image after "+str(args.epoch)+" epoch to local directory named output.png")
